#!/usr/bin/env python
# -*- coding:utf-8 -*-


'''
@CreateTime: 2020/10/27 10:47
@Author : shouke
'''


import socket
from uuid import uuid4
from rpc.protocol import Message
from rpc.zmqrpc import Client
from config.config import slave_bind_host, master_bind_port
# from scenarios import *
import importlib
from utils.utils import Utils

def close_client(client, no_wait=False):
    if no_wait:
        client.send(Message("quit", None, client.id))
        client.close()
        return

    while True:
        try:
            msg = client.recv()
            if msg.data.get('success'):
                print('操作成功：%s' % msg.data.get('msg'))
                break
            else:
                print('操作失败：%s' % msg.data.get('msg'))
                break
        except Exception as e:
            print("操作失败: %s" % (e))
            break
    client.send(Message("quit", None, client.id))
    client.close()


def print_scenario_config_files():
    '''打印场景配置文件列表'''

    global scenario_config_file_list
    global scenario_config_file_index_list

    scenario_config_file_index_list = []
    scenario_config_file_list = Utils.get_scenario_config_files('./scenarios')
    for index, value in enumerate(scenario_config_file_list):
        print(index + 1, value, end= ' '*4)
        if not (index+1) % 5:
            print()
        scenario_config_file_index_list.append(index + 1)
    print('\n')


if __name__ == '__main__':
    scenario_config = None # 存放场景配置
    scenario_config_file_list = [] # 场景配置文件列表
    scenario_config_file_index_list = []  # 场景配置文件索引列表
    choice = 0
    while choice not in scenario_config_file_index_list:
        print_scenario_config_files()
        print('请输入上述数字编号，选择待压测场景配置：', end='')
        choice = input()
        if not choice.isdigit():
            print('输入有误，请输入上述数字编号\n')
            continue
        choice = int(choice)
        if choice not in scenario_config_file_index_list:
            continue
        else:
            print('已选择场景配置为：%s' % scenario_config_file_list[choice-1])
            scenario = importlib.import_module(scenario_config_file_list[choice-1].lstrip('./').rstrip('.py').replace('/', '.'))
            scenario_config = scenario.scenario_config

    print()
    choice = 0
    client_id = socket.gethostname() + "_" + uuid4().hex
    while choice not in (1, 2, 3, 4, 8, 9):
        print('1-启动压测', '2-停止压测', '3-调整启动速率', '4-调整停止速率', '8-关闭压测程序', '9-退出controller')
        print('请根据实际需要，输入上述数字编号：', end='')
        choice = input()
        if not choice.isdigit():
            print('输入有误，请输入上述数字编号\n')
            continue
        choice = int(choice)
        if choice not in (1, 2, 3, 4, 8, 9):
            continue
        else:
            if choice == 1:
                client = Client(slave_bind_host, master_bind_port, client_id)
                if scenario_config:
                    client.send(Message("runLoadTest", scenario_config, client_id))
                    close_client(client)
                else:
                    close_client(client, True)
            elif choice == 2:
                client = Client(slave_bind_host, master_bind_port, client_id)
                client.send(Message("stopLoadTest", None, client_id))
                close_client(client)
            elif choice == 3:
                client = Client(slave_bind_host, master_bind_port, client_id)
                client.send(Message("resetSpawnRate", {"spawn_rate": 20}, client_id))
                close_client(client)
            elif choice == 4:
                client = Client(slave_bind_host, master_bind_port, client_id)
                client.send(Message("resetStopRate", {"stop_rate": 20}, client_id))
                close_client(client)
            elif choice == 8:
                print('正在关闭压测程序')
                client = Client(slave_bind_host, master_bind_port, client_id)
                client.send(Message("quitMaster", None, client_id))
                client.close()
            elif choice == 9:
                exit()
            choice = 0